package com.ruoyi.common.utils;

import com.alibaba.fastjson2.JSONObject;
import com.ruoyi.common.config.LlmInfoConfig;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

/**
 * dify 工作流对话流调用工具
 */
@Slf4j
public class DIfyUtils {

    /**
     * 问答工作流 流式输出
     *
     * @param Auth       token
     * @param jsonObject 参数
     * @return 输出流
     */
    public static SseEmitter runWorkFlowStream(JSONObject jsonObject, String Auth) {
        // 设置超时时间 300s
        SseEmitter emitter = new SseEmitter(300_000L);

        // 使用异步线程处理流式响应
        ExecutorService executor = Executors.newSingleThreadExecutor();
        executor.execute(() -> {
            try {
                String runUrl = LlmInfoConfig.getBaseUrl() + ":81/v1/chat-messages";
                log.info("runUrl: {}", runUrl);

                // 创建请求头
                Map<String, String> headerMap = new HashMap<>();
                headerMap.put("Authorization", "Bearer " + Auth);
                headerMap.put(HttpHeaders.ACCEPT, MediaType.TEXT_EVENT_STREAM_VALUE);
                // 使用restTemplate的execute方法处理流式响应
                RestTemplateClient.post(
                        runUrl,
                        headerMap,
                        jsonObject.toJSONString(),
                        response -> {
                            try (BufferedReader reader = new BufferedReader(new InputStreamReader(response.getBody()))) {
                                String line;
                                while ((line = reader.readLine()) != null) {
                                    log.info("line: {}", line);
                                    // 将每一行数据推送到前端
                                    emitter.send(SseEmitter.event().data(line));
                                }
                                // 完成推送
                                emitter.complete();
                            } catch (Exception e) {
                                emitter.completeWithError(e);
                            }
                            return null;
                        }
                );

            } catch (Exception e) {
                log.error("处理过程中发生错误: {}", e.getMessage());
                emitter.completeWithError(e);
            } finally {
                log.info("流式输出结束...");
            }
        });
        // 关闭线程池
        executor.shutdown();
        log.info("流式输出完成...");
        return emitter;
    }


}
